/******************************************************************************
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 * libkperf licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 * Author:
 * Create: 2025-01-13
 * Description: Utility function implementation for thread binding
 ******************************************************************************/

#include "../include/utils.h"
#include <unistd.h>
#include <sys/syscall.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <cstdio>
#include <cstring>
#include <cerrno>
#include <ctime>
#include <algorithm>
#include <signal.h>
#include <thread>

namespace ThreadBinding {
namespace Utils {

SystemInfo get_system_info() {
    SystemInfo info;
    // 读取/proc/cpuinfo和/sys/devices/system/node/ 以获取NUMA和CPU拓扑
    // 这里只做简单实现，生产环境建议用hwloc等库
    int max_cpu = sysconf(_SC_NPROCESSORS_CONF);
    info.total_cpus = max_cpu;
    info.cpu_to_numa_map.resize(max_cpu, 0);
    info.numa_nodes = 0;
    // 读取/sys/devices/system/node/node*/cpulist
    for (int node = 0; ; ++node) {
        std::string path = "/sys/devices/system/node/node" + std::to_string(node) + "/cpulist";
        std::ifstream in(path);
        if (!in.is_open()) break;
        std::string line;
        getline(in, line);
        std::vector<int> cpus = parse_cpu_range(line);
        if (!cpus.empty()) {
            info.numa_to_cpus_map[node] = cpus;
            for (int cpu : cpus) {
                if (cpu < (int)info.cpu_to_numa_map.size())
                    info.cpu_to_numa_map[cpu] = node;
            }
            info.numa_nodes++;
        }
    }
    return info;
}

CpuAffinityInfo get_thread_cpu_affinity(int tid) {
    CpuAffinityInfo info;
    std::string path = "/proc/" + std::to_string(tid) + "/status";
    std::ifstream in(path);
    if (!in.is_open()) return info;
    std::string line;
    const std::string key = "Cpus_allowed_list:";
    while (getline(in, line)) {
        if (line.find(key) == 0) {
            size_t pos = line.find_first_not_of(" \t", key.size());
            std::string range = (pos != std::string::npos) ? line.substr(pos) : "";
            info.allowed_cpus = parse_cpu_range(range);
            if (!info.allowed_cpus.empty()) {
                info.is_valid = true;
                // 取第一个CPU的NUMA作为主NUMA
                info.numa_id = get_numa_id_for_core(info.allowed_cpus[0]);
            }
            break;
        }
    }
    return info;
}

CpuAffinityInfo get_process_cpu_affinity(int pid) {
    return get_thread_cpu_affinity(pid); // 进程主线程tid=pid
}

bool bind_thread_to_core(int tid, int core_id) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(core_id, &cpuset);
    int ret = sched_setaffinity(tid, sizeof(cpu_set_t), &cpuset);
    return ret == 0;
}

std::vector<int> parse_cpu_range(const std::string& range_str) {
    std::vector<int> cpus;
    std::stringstream ss(range_str);
    std::string part;
    while (getline(ss, part, ',')) {
        size_t hyphen = part.find('-');
        if (hyphen != std::string::npos) {
            int start = std::stoi(part.substr(0, hyphen));
            int end = std::stoi(part.substr(hyphen + 1));
            for (int i = start; i <= end; ++i) cpus.push_back(i);
        } else {
            try { cpus.push_back(std::stoi(part)); } catch (...) {}
        }
    }
    std::sort(cpus.begin(), cpus.end());
    cpus.erase(std::unique(cpus.begin(), cpus.end()), cpus.end());
    return cpus;
}

int get_numa_id_for_core(int core_id) {
    std::string path = "/sys/devices/system/cpu/cpu" + std::to_string(core_id) + "/node0";
    struct stat st;
    for (int node = 0; node < 8; ++node) { // 假定最多8个NUMA
        path = "/sys/devices/system/cpu/cpu" + std::to_string(core_id) + "/node" + std::to_string(node);
        if (stat(path.c_str(), &st) == 0) return node;
    }
    return -1;
}

std::vector<int> find_processes_by_keyword(const std::string& keyword) {
    std::vector<int> pids;
    DIR* dir = opendir("/proc");
    if (!dir) return pids;
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        if (entry->d_type == DT_DIR) {
            char* endptr;
            long pid = strtol(entry->d_name, &endptr, 10);
            if (*endptr == '\0') {
                std::string cmdline_path = "/proc/" + std::string(entry->d_name) + "/cmdline";
                std::ifstream cmdline(cmdline_path);
                if (cmdline.is_open()) {
                    std::string cmd;
                    getline(cmdline, cmd);
                    if (cmd.find(keyword) != std::string::npos) {
                        pids.push_back(pid);
                    }
                }
            }
        }
    }
    closedir(dir);
    return pids;
}

std::vector<int> get_process_threads(int pid) {
    std::vector<int> tids;
    std::string task_path = "/proc/" + std::to_string(pid) + "/task";
    DIR* dir = opendir(task_path.c_str());
    if (!dir) return tids;
    
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        if (entry->d_type == DT_DIR) {
            char* endptr;
            long tid = strtol(entry->d_name, &endptr, 10);
            if (*endptr == '\0' && tid > 0) {
                tids.push_back(tid);
            }
        }
    }
    closedir(dir);
    
    // 如果没有找到线程，至少包含主线程（进程ID）
    if (tids.empty() && process_exists(pid)) {
        tids.push_back(pid);
    }
    
    return tids;
}

std::string get_process_name(int pid) {
    std::string path = "/proc/" + std::to_string(pid) + "/comm";
    std::ifstream in(path);
    std::string name;
    if (in.is_open()) getline(in, name);
    return name;
}

std::string get_thread_name(int tid) {
    return get_process_name(tid); // 简化实现
}

bool process_exists(int pid) {
    return (kill(pid, 0) == 0);
}

bool thread_exists(int tid) {
    return process_exists(tid);
}

std::string format_duration(uint64_t duration_ms) {
    std::ostringstream oss;
    oss << duration_ms << " ms";
    return oss.str();
}

std::string format_percentage(double value, double total, int precision) {
    std::ostringstream oss;
    if (total == 0) oss << "0%";
    else {
        oss << std::fixed << std::setprecision(precision) << (value / total * 100.0) << "%";
    }
    return oss.str();
}

uint64_t get_current_timestamp_ms() {
    return static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(
        std::chrono::steady_clock::now().time_since_epoch()).count());
}

void sleep_ms(uint64_t ms) {
    std::this_thread::sleep_for(std::chrono::milliseconds(ms));
}

void log_message(const std::string& message, bool verbose) {
    if (verbose) {
        std::cout << "[INFO] " << message << std::endl;
    }
}

void log_error(const std::string& message) {
    std::cerr << "[ERROR] " << message << std::endl;
}

void log_warning(const std::string& message) {
    std::cerr << "[WARN] " << message << std::endl;
}

void log_debug(const std::string& message) {
    std::cout << "[DEBUG] " << message << std::endl;
}

} // namespace Utils
} // namespace ThreadBinding 